#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

void translate_main(CodeGen *codegen); // 将 main 函数翻译为汇编代码

int main() {
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 告诉汇编器将汇编放到代码段
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    translate_main(codegen);

    std::cout << codegen->print();
    delete codegen;
    delete module;
    return 0;
}

// TODO: 按照自己的理解翻译main函数，可以按照注释中的步骤，也可自由发挥
void translate_main(CodeGen *codegen) {
    std::unordered_map<std::string, int> offset_map;

    /* Declare main function */
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute);

    /* main function start */
    codegen->append_inst("main", ASMInstruction::Label);

    /* Prologue */
    codegen->append_inst("addi sp, sp, -48");
    codegen->append_inst("sd ra, 40(sp)");
    codegen->append_inst("sd s0, 32(sp)");
    codegen->append_inst("addi s0, sp, 48");

    /* label_entry */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);

    /* %op0 = icmp sgt i32 5, 1 */
    codegen->append_inst("%op0 = icmp sgt i32 5, 1", ASMInstruction::Comment);
    offset_map["%op0"] = -24;
    codegen->append_inst("li", {"t0", "5"});
    codegen->append_inst("li", {"t1", "1"});
    codegen->append_inst("sgt", {"t0", "t0", "t1"}); // t0 = (5 > 1) ? 1 : 0
    codegen->append_inst("sd", {"t0", "s0", "-24"});

    /* %op1 = zext i1 %op0 to i32 */
    codegen->append_inst("%op1 = zext i1 %op0 to i32", ASMInstruction::Comment);
    offset_map["%op1"] = -32;
    codegen->append_inst("ld", {"t0", "s0", "-24"});
    codegen->append_inst("sw", {"t0", "s0", "-32"}); // i1 is 0/1, store as i32

    /* %op2 = icmp ne i32 %op1, 0 */
    codegen->append_inst("%op2 = icmp ne i32 %op1, 0", ASMInstruction::Comment);
    offset_map["%op2"] = -40;
    codegen->append_inst("lw", {"t0", "s0", "-32"});
    codegen->append_inst("snez", {"t0", "t0"}); // t0 = (t0 != 0) ? 1 : 0
    codegen->append_inst("sd", {"t0", "s0", "-40"});

    /* br i1 %op2, label %label3, label %label4 */
    codegen->append_inst("br i1 %op2, label %label3, label %label4", ASMInstruction::Comment);
    codegen->append_inst("ld", {"t0", "s0", "-40"});
    codegen->append_inst("bnez", {"t0", ".main_label3"});
    codegen->append_inst("j", {".main_label4"});

    /* label3: */
    codegen->append_inst(".main_label3", ASMInstruction::Label);

    /* ret i32 233 */
    codegen->append_inst("ret i32 233", ASMInstruction::Comment);
    codegen->append_inst("li", {"a0", "233"});
    codegen->append_inst("j main_exit");

    /* label4: */
    codegen->append_inst(".main_label4", ASMInstruction::Label);

    /* ret i32 0 */
    codegen->append_inst("ret i32 0", ASMInstruction::Comment);
    codegen->append_inst("li", {"a0", "0"});
    codegen->append_inst("j main_exit");

    /* Epilogue */
    codegen->append_inst("main_exit", ASMInstruction::Label);
    codegen->append_inst("ld ra, 40(sp)");
    codegen->append_inst("ld s0, 32(sp)");
    codegen->append_inst("addi sp, sp, 48");
    codegen->append_inst("jr ra");
}